# __init__.py
import os
import sys
from pathlib import Path

# 使用传统方式，但更加健壮
PROJECT_ROOT_STR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT_STR not in sys.path:
    sys.path.insert(0, PROJECT_ROOT_STR)
    
from config.path_config import *
"""PPQ量化主程序入口"""
from config.quantization_config import create_default_config
from core.quantization_core import QuantizationPipeline

def main(): 
    """主函数 - 演示如何使用重构后的代码"""
    # 创建配置
    config = create_default_config()
    
    # 可以自定义配置
    config.calib_steps = 64
    config.batch_size = 1

    # 创建量化流水线
    pipeline = QuantizationPipeline(config)
    
    # 运行完整流水线
    try:
        quantized_model = pipeline.run_complete_pipeline(enable_advanced_optimization=True)
        print("量化完成！")
    except Exception as e:
        print(f"量化过程出错: {e}")
        raise
    

def run_custom_pipeline():
    """自定义流水线示例"""
    config = create_default_config()
        # 可以自定义配置
    config.calib_steps = 6
    config.batch_size = 1
    
    pipeline = QuantizationPipeline(config)
    
    print("运行自定义量化流水线...")
    
    # 1. 设置量化参数
    print("步骤 1/10: 设置量化参数")
    pipeline.setup_quantization_settings()
    
    # 2. 加载模型
    print("步骤 2/10: 加载ONNX模型")
    pipeline.load_model()
    
    # 3. 验证输入
    print("步骤 3/10: 验证模型输入")
    pipeline.validate_graph_inputs()  
    
    # 4. 设置Split-Concat路径为FP32
    print("步骤 4/10: 设置Split-Concat路径")
    pipeline.assign_split_concat_path_to_fp32()
    
    # 5. 构建校准数据集
    print("步骤 5/10: 构建校准数据集")
    pipeline.setup_calibration_dataset()
    
    # 6. 执行基础量化
    print("步骤 6/10: 执行基础量化")
    pipeline.perform_basic_quantization()

    # 7.. 'percentile' 训练
    pipeline.settings_manager.set_calibration_algorithm('percentile')
    
    # 8. 优化 + 量化
    pipeline.apply_advanced_optimization()
    
    # 9. 导出
    pipeline.export_quantized_model()
    print("自定义量化流水线完成！")

def run_basic_quantization():
    """基础量化流水线（无高级优化）"""
    config = create_default_config()
    # 禁用高级优化相关设置
    config.lsq_optimization = False
    
    pipeline = QuantizationPipeline(config)
    
    # 运行基础流水线
    quantized_model = pipeline.run_complete_pipeline(enable_advanced_optimization=False)
    print("基础量化完成！")

if __name__ == "__main__":
    # main()
    run_custom_pipeline()

